"""
矩阵运算工具

提供分布式优化中常用的矩阵运算
"""

import numpy as np
from typing import List, Tuple, Optional
import scipy.linalg as la


def generate_doubly_stochastic_matrix(adjacency: np.ndarray,
                                     method: str = 'metropolis') -> np.ndarray:
    """
    生成双随机矩阵
    
    参数:
        adjacency: 邻接矩阵
        method: 生成方法 ('metropolis', 'laplacian', 'max-degree')
        
    返回:
        双随机权重矩阵
    """
    n = adjacency.shape[0]
    W = np.zeros((n, n))
    
    if method == 'metropolis':
        # Metropolis-Hastings权重
        degrees = np.sum(adjacency, axis=1)
        
        for i in range(n):
            for j in range(n):
                if i != j and adjacency[i, j] > 0:
                    W[i, j] = 1 / (1 + max(degrees[i], degrees[j]))
                    
        # 对角元素
        for i in range(n):
            W[i, i] = 1 - np.sum(W[i, :])
            
    elif method == 'laplacian':
        # 基于拉普拉斯矩阵
        degrees = np.sum(adjacency, axis=1)
        D = np.diag(degrees)
        L = D - adjacency
        
        # 归一化拉普拉斯
        epsilon = 0.5
        W = np.eye(n) - epsilon * L / np.max(degrees)
        
    elif method == 'max-degree':
        # 最大度方法
        degrees = np.sum(adjacency, axis=1)
        max_degree = np.max(degrees)
        
        W = adjacency / max_degree
        for i in range(n):
            W[i, i] = 1 - degrees[i] / max_degree
            
    else:
        raise ValueError(f"未知的方法: {method}")
        
    return W
    

def check_doubly_stochastic(W: np.ndarray, tol: float = 1e-10) -> bool:
    """检查矩阵是否为双随机矩阵"""
    # 检查行和
    row_sums = np.sum(W, axis=1)
    if not np.allclose(row_sums, 1, atol=tol):
        return False
        
    # 检查列和
    col_sums = np.sum(W, axis=0)
    if not np.allclose(col_sums, 1, atol=tol):
        return False
        
    # 检查非负性
    if np.any(W < -tol):
        return False
        
    return True
    

def compute_consensus_error(values: Dict[int, np.ndarray]) -> float:
    """
    计算共识误差
    
    参数:
        values: 节点值字典
        
    返回:
        共识误差 ||x_i - x_avg||_F
    """
    # 计算平均值
    num_nodes = len(values)
    avg_value = np.zeros_like(next(iter(values.values())))
    
    for value in values.values():
        avg_value += value
    avg_value /= num_nodes
    
    # 计算误差
    error = 0.0
    for value in values.values():
        error += np.linalg.norm(value - avg_value)**2
        
    return np.sqrt(error)
    

def compute_gradient_tracking_error(gradients: Dict[int, np.ndarray],
                                   y_values: Dict[int, np.ndarray]) -> float:
    """计算梯度跟踪误差"""
    num_nodes = len(gradients)
    
    # 计算平均梯度
    avg_grad = np.zeros_like(next(iter(gradients.values())))
    for grad in gradients.values():
        avg_grad += grad
    avg_grad /= num_nodes
    
    # 计算误差
    error = 0.0
    for i, y in y_values.items():
        if i in gradients:
            error += np.linalg.norm(y - gradients[i] + avg_grad)**2
            
    return np.sqrt(error)
    

def power_iteration(W: np.ndarray, max_iter: int = 1000,
                   tol: float = 1e-10) -> Tuple[float, np.ndarray]:
    """
    幂迭代法计算最大特征值和特征向量
    
    参数:
        W: 权重矩阵
        max_iter: 最大迭代次数
        tol: 收敛容差
        
    返回:
        (最大特征值, 对应特征向量)
    """
    n = W.shape[0]
    v = np.random.randn(n)
    v = v / np.linalg.norm(v)
    
    lambda_old = 0
    
    for _ in range(max_iter):
        v_new = W @ v
        lambda_new = np.linalg.norm(v_new)
        v_new = v_new / lambda_new
        
        if abs(lambda_new - lambda_old) < tol:
            break
            
        v = v_new
        lambda_old = lambda_new
        
    return lambda_new, v
    

def compute_mixing_time(W: np.ndarray, epsilon: float = 0.01) -> int:
    """
    计算混合时间
    
    参数:
        W: 双随机矩阵
        epsilon: 目标精度
        
    返回:
        达到epsilon-混合所需的迭代次数
    """
    eigenvalues = np.linalg.eigvals(W)
    eigenvalues = np.sort(np.abs(eigenvalues))[::-1]
    
    # 第二大特征值
    if len(eigenvalues) > 1:
        rho = eigenvalues[1]
    else:
        rho = 0
        
    if rho == 0:
        return 1
        
    # 混合时间估计
    mixing_time = int(np.ceil(np.log(epsilon) / np.log(rho)))
    
    return max(1, mixing_time)
    

def kronecker_product(A: np.ndarray, B: np.ndarray) -> np.ndarray:
    """计算Kronecker积"""
    return np.kron(A, B)
    

def vec(A: np.ndarray) -> np.ndarray:
    """矩阵向量化（按列堆叠）"""
    return A.flatten('F')
    

def unvec(v: np.ndarray, shape: Tuple[int, int]) -> np.ndarray:
    """向量反向量化"""
    return v.reshape(shape, order='F')
    

def compute_effective_resistance(L: np.ndarray) -> np.ndarray:
    """
    计算图的有效电阻矩阵
    
    参数:
        L: 拉普拉斯矩阵
        
    返回:
        有效电阻矩阵
    """
    n = L.shape[0]
    
    # 计算Moore-Penrose伪逆
    L_pinv = np.linalg.pinv(L)
    
    # 有效电阻矩阵
    R = np.zeros((n, n))
    
    for i in range(n):
        for j in range(i+1, n):
            e_ij = np.zeros(n)
            e_ij[i] = 1
            e_ij[j] = -1
            
            R[i, j] = e_ij.T @ L_pinv @ e_ij
            R[j, i] = R[i, j]
            
    return R 